package org.jeeasy.generate.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.ServletResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.jeeasy.common.core.domain.vo.R;
import org.jeeasy.generate.domain.dto.CreateGenTableByAiDTO;
import org.jeeasy.generate.domain.vo.AiGenTableVO;
import org.springframework.ai.chat.messages.AbstractMessage;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.content.Media;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.MediaType;
import org.springframework.util.CollectionUtils;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@RestController
@Tag(name = "代码生成")
@RequestMapping("/gen/ai")
@RequiredArgsConstructor
public class AiGeneratorController {

    //final GeneratorService generatorService;

    //@Lazy
    //@Resource
    //@Qualifier("OllamaService")
    //IAiChatApiFeign aiChatService;

    final OpenAiChatModel openaiChatModel;
    final OllamaChatModel ollamaChatModel;
    //final ChatClient chatClient;


    @PostMapping(value = "/create")
    @Operation(summary = "执行")
    public R<AiGenTableVO> generator(@RequestBody CreateGenTableByAiDTO module, ServletResponse response) {
        log.info("module:{}", module);
        //List<Media> mediaList = Arrays.stream(module.getImages()).map(image -> new Media(MimeTypeUtils.IMAGE_PNG, new ByteArrayResource(Base64.getDecoder().decode(image)))).toList();
        //response.setCharacterEncoding("UTF-8");
        Message systemMessage = new SystemPromptTemplate("""
                # 角色
                你是一位经验丰富的数据库管理员，同时也是以为资深程序员，专长在于根据业务需求设计高效且合理的数据库结构和简洁高性能的代码。
                
                ## 任务要求
                1. **设计图解读**: 如果用户给出设计图片，根据设计图生成相关字段
                2. **字段补全**: 深入分析功能需求，如果用户列举的字段不足以满足完成需求，尽可能的补全字段。
                
                ## 限制
                1. **忽略公共字段**: id、创建时间、创建人、修改时间、修改人、删除标记、版本号、启用标记等
                2. **表名、字段名**: 使用字母小写+下划线"_"
                3. **描述字段**: 尽可能使用中文
                """).createMessage();

        BeanOutputConverter<AiGenTableVO> outputConverter = new BeanOutputConverter<>(AiGenTableVO.class);

        String format = outputConverter.getFormat();
        log.debug("outputConverter.getFormat() = {}", format);
        String template = """
                根据下面的描述完成设计:
                {description}.
                {aiDescription}
                {format}
                """;
        String columnDescription = getColumnDescription(module);
        log.debug("columnDescription = {}", columnDescription);
        //PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format, "description", module.getDescription(), "aiDescription", columnDescription));
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("format", format, "description", module.getDescription(), "aiDescription", columnDescription))
                .build();

        Message userMessage = promptTemplate.createMessage();
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage), ChatOptions.builder().model("hf:Qwen/Qwen2.5-Coder-32B-Instruct").build());
        //Prompt prompt = new Prompt(List.of(systemMessage, userMessage),
        //        OllamaOptions.builder().model(module.getModel())
        //                //.withTemperature(0.4)
        //                .build());

        //ChatResponse response = this.openaiChatModel.call(prompt);
        //Generation generation = response.getResult();


        //String text = generation.getOutput().getText();
        //log.debug("generation.getOutput().getText() = {}", text);
        //AiGenTableVO genTableVO = outputConverter.convert(text);

        //Message userMessage = new UserMessage(module.getDescription(), mediaList);
        //Prompt prompt = new Prompt(List.of(systemMessage, userMessage),
        //        OllamaOptions.builder().model(module.getModel())
        //                //.withTemperature(0.4)
        //                .build());
        //AiGenTableVO genTableVO = ChatClient.create(chatModel).prompt(prompt)
        //        //.user(u -> u.text("Generate the filmography of 5 movies for {actor}.")
        //        //        .param("actor", "Tom Hanks"))
        //        .call()
        //        .entity(AiGenTableVO.class);


        String generationTextFromStream = Objects.requireNonNull(this.openaiChatModel.stream(prompt)
                        .collectList()
                        .block())
                .stream()
                .map(ChatResponse::getResults)
                .flatMap(List::stream)
                .map(Generation::getOutput)
                .map(AssistantMessage::getText)
                .filter(Objects::nonNull)
                .collect(Collectors.joining());

        log.debug("generationTextFromStream = {}", generationTextFromStream);
        AiGenTableVO aiGenTableVO = outputConverter.convert(generationTextFromStream);
        log.debug("aiGenTableVO = {}", aiGenTableVO);

        return R.ok(aiGenTableVO);
        //return openaiChatModel.stream(prompt)
        //        // 过滤掉空的聊天响应对象
        //        .filter(Objects::nonNull)
        //        // 过滤掉结果为空的聊天响应
        //        .filter(chatResponse -> chatResponse.getResults() != null)
        //        // 将每个聊天响应的结果转换为Flux流
        //        .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))
        //        // 过滤掉空的结果项
        //        .filter(Objects::nonNull)
        //        // 提取生成的输出
        //        .map(Generation::getOutput)
        //        // 过滤掉空的输出项
        //        .filter(Objects::nonNull)
        //        // 进一步过滤掉内容为空的输出
        //        .filter(content -> Objects.nonNull(content.getContent()))
        //        // 提取助手消息的内容
        //        .map(AssistantMessage::getContent)
        //        // 再次过滤掉空的内容项
        //        .filter(Objects::nonNull)
        //        // 使用reduce收集所有的内容
        //        //.reduce("", (acc, text) -> {
        //        //    log.info("text:{}", text);
        //        //    return acc + text;
        //        //})
        //        // 在内容项后拼接完成信号"complete"
        //        //.map(content -> content)
        //        // 在内容项后拼接完成信号"complete"
        //        .concatWith(Flux.just("complete"));

        //return R.ok(genTableVO);
    }

    private String getColumnDescription(CreateGenTableByAiDTO module) {
        if (CollectionUtils.isEmpty(Arrays.asList(module.getImages()))) {
            return StringUtils.EMPTY;
        }
        List<Media> mediaList = Arrays.stream(module.getImages()).map(image -> new Media(MimeTypeUtils.IMAGE_PNG, new ByteArrayResource(Base64.getDecoder().decode(image)))).toList();

        Message systemMessage = new SystemPromptTemplate("""
                # 角色
                您是软件开发产品经理，尤其擅长产品需求分解，可以根据UI设计图提取精准的业务属性交由程序员更快的开发业务程序。
                
                ## 技能
                ### 设计图解读，根据用户给出设计图片，提取业务业务属性和字段。
                
                ## 限制
                - 严格按照设计图输出业务字段
                - 忽略id、创建时间、创建人、修改事件、修改人等公共字段
                
                ## 输出格式:
                    功能描述: ***
                
                    字段设计:
                    field_name1: 字段1
                    field_name2: 字段2
                    field_name3: 字段3
                """).createMessage();

        String template = """
                {description}.
                """;
        //PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("description", module.getDescription()));
        PromptTemplate promptTemplate = PromptTemplate.builder()
                .template(template)
                .variables(Map.of("description", module.getDescription()))
                .build();

        //Message userMessage = CollectionUtils.isEmpty(mediaList) ? promptTemplate.createMessage() : promptTemplate.createMessage(mediaList);
        //Prompt prompt = new Prompt(List.of(systemMessage, userMessage), ChatOptions.builder().model("llama3.2-vision:latest").build());
        //Message userMessage = new UserMessage("图片是一个form表单页面，解读图片并提取表单项", mediaList);
        Message userMessage = UserMessage.builder()
                .text("图片是一个form表单页面，解读图片并提取表单项")
                .media(mediaList)
                .build();
        Prompt prompt = new Prompt(List.of(userMessage), ChatOptions.builder().model("llava-llama3:latest").build());

        return ollamaChatModel.call(prompt).getResult().getOutput().getText();
    }


    @PostMapping(value = "/test", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @Operation(summary = "执行")
    public Flux<String> test(@RequestBody CreateGenTableByAiDTO module) {
        log.info("module:{}", module);
        //List<Media> mediaList = Arrays.stream(module.getImages()).map(image -> new Media(MimeTypeUtils.IMAGE_PNG, new ByteArrayResource(Base64.getDecoder().decode(image)))).toList();
        List<Media> mediaList = Arrays.stream(module.getImages()).map(image -> Media.builder().data(image).mimeType(MimeTypeUtils.IMAGE_PNG).build()).toList();

        //Message systemMessage = new SystemPromptTemplate("""
        //        # 角色
        //        你是一位经验丰富的数据库管理员，同时也是以为资深程序员，专长在于根据业务需求设计高效且合理的数据库结构和简洁高性能的代码。
        //
        //        ## 任务要求
        //        1. **设计图解读**: 如果用户给出设计图片，根据设计图生成相关字段
        //        2. **字段补全**: 深入分析功能需求，如果用户列举的字段不足以满足完成需求，尽可能的补全字段。
        //        ## 限制
        //        1. **忽略公共字段**: id、创建时间、创建人、修改时间、修改人、删除标记、版本号、启用标记等
        //        2. **表名、字段名**: 使用字母小写+下划线"_"
        //        """).createMessage();

        BeanOutputConverter<AiGenTableVO> outputConverter = new BeanOutputConverter<>(AiGenTableVO.class);

        String format = outputConverter.getFormat();
        //log.debug("outputConverter.getFormat() = {}", format);
        String template = """
                {description}.
                """;
        //PromptTemplate promptTemplate = new PromptTemplate(template, Map.of("format", format, "description", module.getDescription()));
        PromptTemplate promptTemplate = PromptTemplate.builder()
                //.renderer(StTemplateRenderer.builder().startDelimiterToken('<').endDelimiterToken('>').build())
                .variables(Map.of("format", format, "description", module.getDescription()))
                .template(template)
                .build();

        Message userMessage = promptTemplate.createMessage(mediaList);
        //UserMessage userMessage = new UserMessage(module.getDescription());
        //Prompt prompt = new Prompt(List.of(userMessage), ChatOptions.builder().model("hf:meta-llama/Llama-3.2-11B-Vision-Instruct").build());
        Prompt prompt = new Prompt(List.of(userMessage), ChatOptions.builder().model("llava-llama3:latest").build());

        //Generation generation = this.chatModel.call(prompt).getResult();

        //AiGenTableVO genTableVO = outputConverter.convert(generation.getOutput().getText());

        //AtomicReference<String> json = new AtomicReference<>("");
        // 根据提示流式处理聊天模型，生成响应
        return ollamaChatModel.stream(prompt)
                // 过滤掉空的聊天响应对象
                .filter(Objects::nonNull)
                // 过滤掉结果为空的聊天响应
                .filter(chatResponse -> chatResponse.getResults() != null)
                // 将每个聊天响应的结果转换为Flux流
                .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))
                // 过滤掉空的结果项
                .filter(Objects::nonNull)
                // 提取生成的输出
                .map(Generation::getOutput)
                // 过滤掉空的输出项
                .filter(Objects::nonNull)
                // 进一步过滤掉内容为空的输出
                .filter(content -> Objects.nonNull(content.getText()))
                // 提取助手消息的内容
                .mapNotNull(AbstractMessage::getText)
                // 再次过滤掉空的内容项
                .filter(Objects::nonNull)
                // 使用reduce收集所有的内容
                .reduce("", (acc, text) -> {
                    log.info("text:{}", text);
                    return acc + text;
                })
                // 在内容项后拼接完成信号"complete"
                .map(content -> content)
                // 在内容项后拼接完成信号"complete"
                .concatWith(Flux.just("complete:"));
    }


}
